import librosa
import numpy as np
import opensmile
from pathlib import Path

# 获取项目根目录的绝对路径
def get_project_root():
    """获取项目根目录路径，兼容不同操作系统和部署环境"""
    current_file = Path(__file__).resolve()
    # 从当前文件向上查找到包含model目录的根目录
    project_root = current_file.parent.parent.parent  # src/service -> src -> project_root
    return project_root

class StressDetector:
    def __init__(self):
        # OpenSMILE 3.x 配置 - 使用eGeMAPS特征集，专为情感分析优化
        try:
            self.smile = opensmile.Smile(
                feature_set=opensmile.FeatureSet.eGeMAPSv02,
                feature_level=opensmile.FeatureLevel.Functionals,
            )
            print("OpenSMILE初始化成功")
        except Exception as e:
            print(f"OpenSMILE初始化失败: {str(e)}")
            print("请确保已正确安装OpenSMILE并配置环境变量")
            raise e

    def preprocess_audio(self, path):
        """音频预处理"""
        y, sr = librosa.load(path, sr=16000, mono=True)
        y = y / (np.max(np.abs(y)) + 1e-10)

        # 改进的静音检测
        energy = librosa.feature.rms(y=y)[0]
        threshold = np.percentile(energy, 20)  # 更鲁棒的阈值
        voiced = energy > threshold
        silence_ratio = 1 - np.mean(voiced)
        
        # 计算语音持续时间
        speech_duration = np.sum(voiced) * (len(y) / len(energy)) / sr
        
        return y, sr, silence_ratio, speech_duration

    def extract_opensmile_features(self, path):
        """提取OpenSMILE特征"""
        try:
            df = self.smile.process_file(path)
            # 打印可用的特征名称，帮助调试
            # print("可用特征列表:", list(df.columns))
            
            # 基础特征 - 确保这些特征在eGeMAPS中可用
            features = {
                'f0_var': df['F0semitoneFrom27.5Hz_sma3nz_stddevNorm'].values[0],
                'spectral_flux': df['spectralFlux_sma3_stddevNorm'].values[0],
                'energy_std': df['loudness_sma3_stddevNorm'].values[0]
            }
            
            # 尝试添加增强特征 - 只有在确认特征存在时才添加
            self.try_add_feature(features, df, 'jitter', 'jitterLocal_sma3nz_stddevNorm')
            self.try_add_feature(features, df, 'shimmer', 'shimmerLocaldB_sma3nz_stddevNorm')
            self.try_add_feature(features, df, 'mfcc1_var', 'mfcc1_sma3_stddevNorm')
            self.try_add_feature(features, df, 'alpha_ratio', 'alphaRatio_sma3_stddevNorm')
            self.try_add_feature(features, df, 'mean_f0', 'F0semitoneFrom27.5Hz_sma3nz_meanRisingSlope')
            
            return features
        except Exception as e:
            print(f"OpenSMILE特征提取失败: {str(e)}")
            return None
            
    def try_add_feature(self, features_dict, dataframe, key, column_name):
        """尝试从DataFrame中添加特征到特征字典，如果特征不存在则跳过"""
        try:
            if column_name in dataframe.columns:
                features_dict[key] = dataframe[column_name].values[0]
        except:
            pass  # 如果特征不存在或提取失败，则忽略

    def calc_energy_entropy(self, y):
        """改进的能量熵计算"""
        energy = librosa.feature.rms(y=y)[0]
        energy = energy / (np.sum(energy) + 1e-10)
        energy = np.clip(energy, 1e-10, 1)
        return -np.sum(energy * np.log(energy + 1e-10)) / 5

    def normalize_feature(self, val, max_val, alpha=1.0):
        """避免异常值影响的归一化方法"""
        val = min(max(val, 0), max_val)  # 确保值在0和max_val之间
        return (val / max_val) ** alpha

    def analyze_temporal_dynamics(self, y, sr):
        """分析语音特征的时间动态变化"""
        # 将音频分成多个帧进行分析
        frame_length = int(0.025 * sr)  # 25ms帧
        hop_length = int(0.010 * sr)    # 10ms跳跃
        
        # 提取每帧的能量
        energy_frames = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]
        
        # 计算能量的变化率
        if len(energy_frames) > 1:
            energy_delta = np.diff(energy_frames)
            energy_variability = np.std(energy_delta) / (np.mean(energy_frames) + 1e-10)
        else:
            energy_variability = 0
            
        return energy_variability

    def analyze(self, path):
        """完整分析流程"""
        # 1. 预处理
        y, sr, silence_ratio, speech_duration = self.preprocess_audio(path)

        # 2. 特征提取
        features = self.extract_opensmile_features(path)
        if not features:
            return None

        # 3. 计算衍生特征
        entropy = self.calc_energy_entropy(y)
        energy_variability = self.analyze_temporal_dynamics(y, sr)

        # 4. 优化的压力指数计算 - 基于研究文献的权重调整
        # 使用更简单的特征组合，确保所有特征都可用
        P = 0.35 * self.normalize_feature(features['f0_var'], 0.5) + \
            0.25 * self.normalize_feature(entropy, 5, 0.8) + \
            0.20 * self.normalize_feature(features['spectral_flux'], 0.5, 0.7) + \
            0.10 * self.normalize_feature(energy_variability, 0.3) + \
            0.10 * silence_ratio
            
        # 如果有额外特征，增强计算
        if 'jitter' in features:
            P = 0.9 * P + 0.1 * self.normalize_feature(features['jitter'], 0.4)
        if 'shimmer' in features:
            P = 0.9 * P + 0.1 * self.normalize_feature(features['shimmer'], 0.5)

        # 5. 结果整理
        # 基础特征
        result = {
            "f0_var": float(features['f0_var']),
            "spectral_flux": float(features['spectral_flux']),
            "energy_entropy": float(entropy),
            "silence_ratio": float(silence_ratio),
            "energy_variability": float(energy_variability),
            "speech_duration": float(speech_duration),
        }
        
        # 添加可选特征到结果中
        for feature in ['jitter', 'shimmer', 'mfcc1_var', 'alpha_ratio', 'mean_f0']:
            if feature in features:
                result[feature] = float(features[feature])
                
        # 最终压力指数放在最后
        result["stress"] = float(P)
        return result


# 使用示例
if __name__ == "__main__":
    import requests
    import tempfile
    import uuid
    import os
    
    try:
        analyzer = StressDetector()
        print("压力检测模型初始化成功")
    except Exception as e:
        print(f"压力检测模型初始化失败: {str(e)}")
        exit(1)

    # 测试音频URL - 您可以替换为您的COS音频链接
    audio_url = "https://zx-1343343346.cos.ap-chongqing.myqcloud.com/ai_interview/user_audio/1916016714226561026/1754304843005.webm"
    
    # 本地测试文件路径（备用）
    project_root = get_project_root()
    local_test_file = project_root / "src" / "temp_audio_wav" / "1747197108253_20250514_123148.wav"
    
    temp_file_path = None
    try:
        # 判断是网络链接还是本地路径
        if audio_url.startswith(('http://', 'https://')):
            print(f"正在下载网络音频文件: {audio_url}")
            
            # 下载网络音频文件
            response = requests.get(audio_url, timeout=30)
            response.raise_for_status()
            
            # 从URL中提取文件扩展名
            from urllib.parse import urlparse
            parsed_url = urlparse(audio_url)
            original_filename = os.path.basename(parsed_url.path)
            file_extension = os.path.splitext(original_filename)[1] or '.webm'  # 默认为.webm
            
            # 创建临时文件，保持原始扩展名
            temp_dir = tempfile.gettempdir()
            temp_filename = f"stress_test_audio_{uuid.uuid4().hex}{file_extension}"
            temp_file_path = os.path.join(temp_dir, temp_filename)
            
            # 保存音频数据
            with open(temp_file_path, 'wb') as temp_file:
                temp_file.write(response.content)
            
            print(f"音频文件已下载到: {temp_file_path}")
            
            # 分析下载的音频
            result = analyzer.analyze(temp_file_path)
        else:
            # 使用本地文件路径
            if local_test_file.exists():
                print(f"分析本地音频文件: {local_test_file}")
                result = analyzer.analyze(str(local_test_file))
            else:
                print(f"本地测试文件不存在: {local_test_file}")
                print("请确保本地测试文件存在，或使用网络音频URL")
                result = None

        # 执行分析
        if result:
            print("=== 分析结果 ===")
            for k, v in result.items():
                print(f"{k:15s}: {v:.4f}" if isinstance(v, float) else f"{k:15s}: {v}")
        else:
            print("分析失败，请检查：")
            print("1. OpenSMILE路径配置")
            print("2. 音频文件是否存在")
            print("3. 依赖库版本是否匹配")
            
    except requests.exceptions.RequestException as e:
        print(f"下载音频文件失败: {e}")
    except Exception as e:
        print(f"压力检测分析过程中出错: {e}")
    finally:
        # 清理临时文件
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
                print(f"已清理临时文件: {temp_file_path}")
            except Exception as e:
                print(f"清理临时文件失败: {e}")